#define AES
#define KEYGEN

//2 threads
//thread 1 from_tap => udp-socket
//thread 2 udp-socket => to_tap

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <unistd.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <strings.h>
#include <string.h>
#include <openssl/evp.h>

#define SEND_BUF_SIZE 1
#define RECV_BUF_SIZE 16

#ifdef AES
#define AES_BLOCK_SIZE 128
#define KEY_SIZE_IN_BYTE 16
#define HASH_ROUNDS 4
#endif

int running = 1;

int socket_fd;

int port_num = 8888;

unsigned char send_buf[SEND_BUF_SIZE];
unsigned char *send_ciph_buf;
unsigned char recv_buf[RECV_BUF_SIZE];
unsigned char *recv_ciph_buf;

struct sockaddr_in si_me, si_other;
socklen_t size_of_si_me, size_of_si_other;

FILE *in;
FILE *out;

#ifdef AES
EVP_CIPHER_CTX en, de;
int aes_init(unsigned char *key_data, int key_data_len, unsigned char *salt, EVP_CIPHER_CTX *e_ctx, EVP_CIPHER_CTX *d_ctx);
unsigned char *aes_encrypt(EVP_CIPHER_CTX *e, unsigned char *plaintext, int *len);
unsigned char *aes_decrypt(EVP_CIPHER_CTX *e, unsigned char *ciphertext, int *len);
unsigned int salt[] = {12345, 54321};
unsigned char *key_data;
int key_data_len;
int *rlen;
int *slen;
#endif

//Call and Check for Error
void cce(int error,char *description)
{
   if(error)
   {
     printf("An error occured while %s!\n",description);
     running = 0;
   }
}

void init()
{
  socket_fd = socket(AF_INET,SOCK_DGRAM,IPPROTO_UDP);
  cce(socket_fd==-1,"creating a socket");
  
  size_of_si_me = sizeof(si_me);
  size_of_si_other = sizeof(si_other);
  
  #ifdef CLIENT
  bzero((char*) &si_me, size_of_si_me);
  
  si_other.sin_family = AF_INET;
  si_other.sin_addr.s_addr = inet_addr("127.0.0.1");
  si_other.sin_port = htons(port_num);
  #endif
  
  #ifdef SERVER
  bzero((char*) &si_me, size_of_si_me);
  
  si_me.sin_family = AF_INET;
  si_me.sin_addr.s_addr = htonl(INADDR_ANY);
  si_me.sin_port = htons(port_num);
  
  cce(bind(socket_fd, (struct sockaddr *) &si_me, size_of_si_me)==-1,"binding");
  #endif
  
  in = stdin;//fopen("testfile","r");
  out = stdout;//fopen("testout","w");
  
  #ifdef AES
  cce(aes_init(key_data, key_data_len, (unsigned char *)&salt, &en, &de)!=0,"aes-initialization");
  rlen = malloc(sizeof(int));
  slen = malloc(sizeof(int));
  key_data = (unsigned char *) "0123456789012345";
  key_data_len = KEY_SIZE_IN_BYTE;
  #endif
}

void *from_tap_to_udp(void *arg)
{
  int cnt;
  while(running)
   {
      cnt = fread(send_buf,sizeof(unsigned char),SEND_BUF_SIZE,in);
      #ifdef DEBUG
      printf("read %d bytes, ",cnt);
      #endif
      //aes encrypt send_buf with length cnt here
      if(cnt>=1)
      {
	#ifdef AES
	*slen = cnt;
	send_ciph_buf= aes_encrypt(&en, send_buf, slen);
	cce(sendto(socket_fd,send_ciph_buf,*slen,0,(struct sockaddr *)&si_other,size_of_si_other)<0,"sending");
	#ifdef DEBUG
	printf("sent %d bytes, ",*slen);
	#endif
	#endif
	#ifndef AES
	cce(sendto(socket_fd,send_buf,cnt,0,(struct sockaddr *)&si_other,size_of_si_other)<0,"sending");
	#endif
      }
   }
   pthread_exit(0);
}

void *from_udp_to_tap(void *arg)
{
  int cnt;
  while(running)
  {
    cce((cnt=recvfrom(socket_fd,recv_buf,RECV_BUF_SIZE,0,(struct sockaddr *)&si_other,&size_of_si_other))<0,"receiving");
    #ifdef DEBUG
    printf("received %d bytes, ",cnt);
    #endif
    //aes decrypt send_buf with length cnt here
    
    if(cnt>=1)
    {
      #ifdef AES
      *rlen = cnt;
      recv_ciph_buf= aes_decrypt(&de, recv_buf,rlen);
      fwrite(recv_ciph_buf,sizeof(char),*rlen,out);
      #ifdef DEBUG
      printf("wrote %d bytes, ",*rlen);
      #endif
      #endif
      #ifndef AES
      fwrite(recv_buf,sizeof(unsigned char),cnt,out);
      #endif
      fflush(out);
    }
  }  
  pthread_exit(0);
}

int handshake()
{
  unsigned char hello[] = "hello";
  #ifdef CLIENT
  cce(sendto(socket_fd,hello,6,0,(struct sockaddr *)&si_other,size_of_si_other)==-1,"client handshake send");
  cce(recvfrom(socket_fd,recv_buf,RECV_BUF_SIZE,0,(struct sockaddr *)&si_other,&size_of_si_other)==-1,"client handshake receive");
  
  return memcmp(hello,recv_buf,5);
  #endif
  
  #ifdef SERVER
  cce(recvfrom(socket_fd,recv_buf,RECV_BUF_SIZE,0,(struct sockaddr *)&si_other,&size_of_si_other)==-1,"server handshake receive");
  recv_buf[6] = '\0';
  
  if(memcmp(hello,recv_buf,5)==0)
    return sendto(socket_fd,hello,6,0,(struct sockaddr *)&si_other,size_of_si_other)==(-1);
  else
    return 2;
  #endif
}

void thread_management()
{
  pthread_t thread_tap2udp, thread_udp2tap;
  
  cce(pthread_create( &thread_tap2udp, NULL, from_tap_to_udp, NULL )!=0,"creating sending thread");
  
  cce(pthread_create( &thread_udp2tap, NULL, from_udp_to_tap, NULL )!=0,"creating reiceving thread");
  
  //wait until both threads terminate
  pthread_join( thread_tap2udp, NULL );
  pthread_join( thread_udp2tap, NULL );
}

int main()
{
  init();
 
  if(handshake()==0)
    thread_management();
  else
    cce(1,"handshake");
    
  close(socket_fd);
  fclose(in);
  fclose(out);
  
  #ifdef AES
  free(rlen);
  free(slen);
  free(recv_ciph_buf);
  free(send_ciph_buf);
  #endif
  return 0;
}


#ifdef AES

void load_key(unsigned char *key,int len,const char* path)
{
  FILE *key_file;
  key_file = fopen(path,"r");
  fread(key,len,sizeof(unsigned char),key_file);
  fclose(key_file);
}

void save_key(unsigned char *key,int len,const char* path)
{
  FILE *key_file;
  key_file = fopen(path,"w");
  fwrite(key,len,sizeof(unsigned char),key_file);
  fclose(key_file);
}

////////////////////////////////////AES//////////////////////////////////////

//SOURCE: http://saju.net.in/code/misc/openssl_aes.c.txt

/**
 * Create an 256 bit key and IV using the supplied key_data. salt can be added for taste.
 * Fills in the encryption and decryption ctx objects and returns 0 on success
 **/
int aes_init(unsigned char *key_data, int key_data_len, unsigned char *salt, EVP_CIPHER_CTX *e_ctx, 
             EVP_CIPHER_CTX *d_ctx)
{
  #ifdef KEYGEN
  #ifdef SERVER
  int i, nrounds = HASH_ROUNDS;
  #endif
  #endif
  unsigned char key[KEY_SIZE_IN_BYTE], iv[KEY_SIZE_IN_BYTE];
  
  #ifdef KEYGEN
  #ifdef SERVER
  /*
   * Gen key & IV for AES 256 CBC mode. A SHA1 digest is used to hash the supplied key material.
   * nrounds is the number of times the we hash the material. More rounds are more secure but
   * slower.
   */
  i = EVP_BytesToKey(EVP_aes_128_cbc(), EVP_sha1(), salt, key_data, key_data_len, nrounds, key, iv);
  if (i != KEY_SIZE_IN_BYTE) {
    printf("Key size is %d bits - should be %d bits\n", i,KEY_SIZE_IN_BYTE);
    return -1;
  }
  save_key(key,KEY_SIZE_IN_BYTE,"key.aes");
  save_key(iv,KEY_SIZE_IN_BYTE,"iv.aes");
  #endif
  #endif
  #ifdef CLIENT
  load_key(key,KEY_SIZE_IN_BYTE,"key.aes");
  load_key(iv,KEY_SIZE_IN_BYTE,"iv.aes");
  #endif
  #ifndef KEYGEN
  #ifdef SERVER
  load_key(key,KEY_SIZE_IN_BYTE,"key.aes");
  load_key(iv,KEY_SIZE_IN_BYTE,"iv.aes");
  #endif
  #endif
  EVP_CIPHER_CTX_init(e_ctx);
  EVP_EncryptInit_ex(e_ctx, EVP_aes_128_cbc(), NULL, key, iv);
  EVP_CIPHER_CTX_init(d_ctx);
  EVP_DecryptInit_ex(d_ctx, EVP_aes_128_cbc(), NULL, key, iv);

  return 0;
}

/*
 * Encrypt *len bytes of data
 * All data going in & out is considered binary (unsigned char[])
 */
unsigned char *aes_encrypt(EVP_CIPHER_CTX *e, unsigned char *plaintext, int *len)
{
  /* max ciphertext len for a n bytes of plaintext is n + AES_BLOCK_SIZE -1 bytes */
  int c_len = *len + AES_BLOCK_SIZE, f_len = 0;
  unsigned char *ciphertext = malloc(c_len);

  /* allows reusing of 'e' for multiple encryption cycles */
  EVP_EncryptInit_ex(e, NULL, NULL, NULL, NULL);

  /* update ciphertext, c_len is filled with the length of ciphertext generated,
    *len is the size of plaintext in bytes */
  EVP_EncryptUpdate(e, ciphertext, &c_len, plaintext, *len);

  /* update ciphertext with the final remaining bytes */
  EVP_EncryptFinal_ex(e, ciphertext+c_len, &f_len);

  *len = c_len + f_len;
  return ciphertext;
}

/*
 * Decrypt *len bytes of ciphertext
 */
unsigned char *aes_decrypt(EVP_CIPHER_CTX *e, unsigned char *ciphertext, int *len)
{
  /* because we have padding ON, we must allocate an extra cipher block size of memory */
  int p_len = *len, f_len = 0;
  unsigned char *plaintext = malloc(p_len + AES_BLOCK_SIZE);
  
  EVP_DecryptInit_ex(e, NULL, NULL, NULL, NULL);
  EVP_DecryptUpdate(e, plaintext, &p_len, ciphertext, *len);
  EVP_DecryptFinal_ex(e, plaintext+p_len, &f_len);

  *len = p_len + f_len;
  return plaintext;
} 
#endif